import atexit
from dataclasses import fields
from time import perf_counter
from tqdm.auto import tqdm
from transformers import AutoTokenizer
import torch.multiprocessing as mp

from nanovllm.config import Config
from nanovllm.sampling_params import SamplingParams
from nanovllm.engine.sequence import Sequence
from nanovllm.engine.scheduler import Scheduler
from nanovllm.engine.model_runner import ModelRunner


class LLMEngine:
    """
    LLM 负责整个推理流程
    1-模型初始化
    2-请求管理和调度
    3-分布式并行
    4-tokenizer处理
    """

    def __init__(self, model, **kwargs):
        #获取所有的config字段，并筛选
        config_fields = {field.name for field in fields(Config)}
        config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields}
        # 初始化配置对象
        config = Config(model, **config_kwargs)
        # 存储分布式进程和事件
        self.ps = []
        self.events = []
        # 使用spawn方式创建多进程上下文
        """
        spawn 是一种进程启动方法，会创建一个全新的 Python 解释器进程
        这种方式特别适合 PyTorch 等需要使用 CUDA 的场景，可以避免一些共享状态的问题
        相比其他启动方式（如 fork），spawn 启动的进程更加隔离，减少了潜在的冲突
        """
        ctx = mp.get_context("spawn")
        # 启动tensor parallel的worker进程（主进程为0号，worker从1开始）
        # 初始化主进程的 ModelRunner
        """
        主进程负责协调其他进程的工作
        它会创建一个 ModelRunner 实例，负责处理主进程的请求
        其他进程会创建自己的 ModelRunner 实例，负责处理分布式并行的请求
        """
        for i in range(1, config.tensor_parallel_size):
            # 创建一个进程间同步事件，用于主进程和worker进程之间的通信与同步
            # 在分布式模型推理中，不同进程需要协调工作，事件机制提供了简单有效的同步方式
            event = ctx.Event()
            # 创建一个新的worker进程，目标函数是ModelRunner，参数包括配置、进程编号i、同步事件
            process = ctx.Process(target=ModelRunner, args=(config, i, event))
            # 启动该worker进程，让其在后台运行
            process.start()
            self.ps.append(process)
            self.events.append(event)
        # 主进程的ModelRunner实例
        self.model_runner = ModelRunner(config, 0, self.events)
        # 加载分词器
        self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True)
        # 设置终止token id
        config.eos = self.tokenizer.eos_token_id
        # 初始化调度器
        self.scheduler = Scheduler(config)
        # 注册退出时的清理函数
        atexit.register(self.exit)

    def exit(self):
        self.model_runner.call("exit")
        del self.model_runner
        for p in self.ps:
            p.join()

    def add_request(self, prompt: str | list[int], sampling_params: SamplingParams):
        if isinstance(prompt, str):
            prompt = self.tokenizer.encode(prompt)
        seq = Sequence(prompt, sampling_params)
        self.scheduler.add(seq)

    def step(self):
        """
        这意味着一次step可能收集0个、1个或多个已完成的序列
        一次step ≠ 一次请求完成，而是模型处理的一个时间片，可能包含：
        多个请求的预填充阶段处理
        多个请求的解码阶段处理
        部分请求完成，部分请求继续处理
        """
        # 选择一批要处理的序列 ，seqs(选中的序列列表)和is_prefill(是否为预填充阶段)
        """
            # Sequence对象2
            {
              seq_id: 2,
              token_ids: [234, 567, 890, 202, 303, 613],  # 添加了新生成的token 613
              status: SequenceStatus.RUNNING,
              # ... 其他属性 ...
            }
        """
        seqs, is_prefill = self.scheduler.schedule()
        # 如果是预填充阶段，并行处理多个新序列的全部提示词tokens
        # 如果是解码阶段，为每个序列生成下一个token
        token_ids = self.model_runner.call("run", seqs, is_prefill)
        # 后处理阶段 更新序列状态
        """
        将生成的token添加到对应序列中
        检查序列是否应该结束（遇到EOS标记或达到最大token数）
        对已完成的序列，释放其KV缓存资源
        """
        self.scheduler.postprocess(seqs, token_ids)
        # 结果收集：只收集本次处理中已完成的序列的输出
        outputs = [(seq.seq_id, seq.completion_token_ids) for seq in seqs if seq.is_finished]
        # 使用正负号来区分预填充和解码阶段的token数统计
        # 预填充阶段：统计所有序列的token数总和
        # 解码阶段：统计完成的序列数（取负号表示解码阶段）
        num_tokens = sum(len(seq) for seq in seqs) if is_prefill else -len(seqs)
        return outputs, num_tokens

    def is_finished(self):
        return self.scheduler.is_finished()

    def generate(
        self,
        prompts: list[str] | list[list[int]],
        sampling_params: SamplingParams | list[SamplingParams],
        use_tqdm: bool = True,
    ) -> list[str]:
        """

        """
        if use_tqdm:
            # 进度条
            pbar = tqdm(total=len(prompts), desc="Generating", dynamic_ncols=True)
        if not isinstance(sampling_params, list):
            sampling_params = [sampling_params] * len(prompts)
        # 按照序列顺序放到一起，批量添加所有请求，而不是立即处理，为调度器的优化提供了空间
        for prompt, sp in zip(prompts, sampling_params):
            # 添加请求到调度器
            self.add_request(prompt, sp)
        outputs = {}
        prefill_throughput = decode_throughput = 0.#浮点数值
        while not self.is_finished():
            # 使用字典存储输出结果，以 seq_id 为键，便于后续按原始顺序排序
            #计时
            t = perf_counter()
            # 模型处理了一批序列的一个时间片。
            output, num_tokens = self.step()
            if use_tqdm:
                if num_tokens > 0:
                    prefill_throughput = num_tokens / (perf_counter() - t)
                else:
                    decode_throughput = -num_tokens / (perf_counter() - t)
                pbar.set_postfix({
                    "Prefill": f"{int(prefill_throughput)}tok/s",
                    "Decode": f"{int(decode_throughput)}tok/s",
                })
            for seq_id, token_ids in output:
                outputs[seq_id] = token_ids
                if use_tqdm:
                    pbar.update(1)
        outputs = [outputs[seq_id] for seq_id in sorted(outputs.keys())]
        outputs = [{"text": self.tokenizer.decode(token_ids), "token_ids": token_ids} for token_ids in outputs]
        if use_tqdm:
            pbar.close()
        return outputs
